We can directly introduce the BatchMeta as it is now, don't need to compare with previous design

same as above comment

The core logics here is not demonstrating the construction, but usage.

What's the relationship between this demo and demonstrate_batch_meta()

Don't need to compare

I believe we don't need to deliberately demonstrate the field_schema here. We can simplify it by showing how to construct a BatchMeta manually

Functionally repeated

better to provide a __str__ representation because there is no where to know how to access the inner properties.

We should only return the field_schema of THE selected sample?

I have tried this code but the returned batch_meta only has one field_schema. Is that because the field has uniform shape? What if it's nested?

BatchMeta(global_indexes=[8, 9, 10, 11, 12, 13, 14, 15], partition_ids=['demo_partition', 'demo_partition', 'demo_partition', 'demo_partition', 'demo_partition', 'demo_partition', 'demo_partition', 'demo_partition'], field_schema={'input_ids': {'dtype': torch.int64, 'shape': torch.Size([512]), 'is_nested': False, 'is_non_tensor': False}, 'attention_mask': {'dtype': torch.float32, 'shape': torch.Size([512]), 'is_nested': False, 'is_non_tensor': False}, 'nested': {'dtype': torch.float32, 'shape': None, 'is_nested': True, 'is_non_tensor': False, 'per_sample_shapes': [(4, 3), (2, 4), (4, 4), (2, 3), (4, 5), (2, 2), (4, 6), (2, 1)]}}, extra_info={}, custom_meta=[{}, {}, {}, {}, {}, {}, {}, {}], _custom_backend_meta=[{}, {}, {}, {}, {}, {}, {}, {}])

That's correct, but as mentioned earlier, better to support extracting the field schema of the given sample. Now it looks like:

I have found a bug:

    input_ids = torch.randint(0, 1000, (8, 512))
    attention_mask = torch.ones(8, 512)

    nested =  torch.nested.as_nested_tensor(
        [torch.randn(4, 3).int(), torch.randn(2, 4),torch.randn(4, 4), torch.randn(2, 3),torch.randn(4, 5), torch.randn(2, 2),torch.randn(4, 6), torch.randn(2, 1)], layout=torch.strided
    )

    data_batch = TensorDict(
        {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "nested": nested,
        },
        batch_size=8,
    )

    partition_id = "demo_partition"
    batch_meta = tq_client.put(data=data_batch, partition_id=partition_id)

The dtype is not supporting per_sample_dtypes and directly use the first element's dtype.

Please check whether the current design fits the backend's requirements @dpj135 @tianyi-ge

Don't need to compare with old version. It will confusing users

Same as above comment

For nested tensor, we cannot assume the first item's dtype equals to all other tensors

Lack per field dtype logic

We should only goes in this loop when the field is nested to speed up?

We'd better remove these deepcopy from __post_init__ since it will bring much higher pressure for controller -- now the controller is called for each single PUT/GET operation.

I suggest to provide a dedicate method that deepcopy these variables, and actively call this method for operations like chunk/union/reorder...

Do we still need local keys now?

It can directly be the global_indexes

unnecessary codes and comments here

There are too many input params. Many of them can be simplified. E.g. chunk_size, num_units

And I suggest to provide a mapping function as earlier version, as it helps to make the dispatching logic more clear.

Why? We don't need _custom_backend_meta to do routing

Can be simplified by providing a mapping function

fully spell as global_indexes. Use gi only for cases like [gi for gi in global_indexes]

Same as above comment

Is there any serial performance comparison for BatchMeta? I'm wondering if pickle could be faster for BatchMeta obj

Why we have to tackle this outside serial_utils.py?

What will happen if:

1. We input a uniform tensor (same dtype & shape) in field A -> only has one dtype and shape info
2. We append other format tensor with different dtype & shape in the same field A? Will we update the dtype and shape info to reflect the new changes?

A folloing question:

If user retrieve the data put in step 1, what's the BatchMeta looks like? It will say this is a nested or ordinary tensor?

per_sample_shapes is only meaningful when is_nested=True, but the validation logic doesn't check this correlation.

Should record a warning:
logger.warning(f"Unknown dtype string '{dtype_str}', returning as-is")

It's better to use English instead of Chinese